// Copyright 2022 jeff.li. and/or its affiliates.

#pragma once

#include <memory>
#include <string>
#include <thread>
#include <vector>

#ifdef _WIN32
#include <windows.h>
typedef DWORD pid_t;
#else

#include <unistd.h>

#endif

#include <tbir/runtime/threadpool/i_thread_pool.h>
#include <tbir/runtime/threadpool/mpmc_bounded_queue.h>

namespace tbir::runtime::internal {

    class LockFreeRunnable : public IRunnable {
    public:
        bool Done() override {
            return finish_;
        }

    protected:
        void SetDone() override {
            finish_ = true;
        }

    private:
        volatile bool finish_ = false;

        friend class LockFreeThreadPool;
    };

    class LockFreeThreadPool : public IThreadPool {
    public:
        explicit LockFreeThreadPool(size_t threads, const std::string &name, int64_t intervals_ns);

        explicit LockFreeThreadPool(size_t threads, const std::string &name);

        ~LockFreeThreadPool() override;

        void Enqueue(IRunnablePtr &runner, size_t seq) override;

        void EnqueueBulk(std::vector<IRunnablePtr> &runners) override;

        size_t GetThreadsNum() const override;

        std::vector<std::thread::id> GetThreadIds() const override;

    protected:
        static void ThreadEntry(LockFreeThreadPool *pool, const std::string &name);

    private:
        // need to keep track of threads so we can join them
        std::vector<std::thread> workers_;
        // the task queue
        MPMCBoundedQueue<IRunnablePtr> tasks_;
        // stop flag
        bool stop_ = false;
        std::string name_;
        int64_t intervals_ns_;
        pid_t belong_to_pid_;
    };

    class SPSCLockFreeThreadPool : public IThreadPool {
    public:
        explicit SPSCLockFreeThreadPool(size_t threads,
                                        const std::string &name,
                                        int64_t intervals_ns = 1);

        ~SPSCLockFreeThreadPool() override;

        void Enqueue(IRunnablePtr &runner, size_t seq) override;

        size_t GetThreadsNum() const override;

        std::vector<std::thread::id> GetThreadIds() const override;

    private:
        std::vector<std::unique_ptr<LockFreeThreadPool>> workers_;
    };

}  // namespace tbir::runtime::internal
